#include "gtest/gtest.h"
#include <mockcpp/mokc.h>
#include <mockcpp/mockcpp.hpp>
 
#include <vector>
#include <tuple>
#include <iostream>
 
#include "topoinfo_struct.h"
#include "log.h"
#include "checker.h"
#include "virtual_topo.h"
#include "testcase_utils.h"
#include "coll_native_executor_base.h"
 
using namespace hccl;

class BroadCastRingFor91093GeneralizationTest: public ::testing::TestWithParam<std::tuple<u64, HcclDataType, OpMode, vector<int>, int,
    std::vector<std::vector<std::vector<unsigned int>>> > > {

public:
    static void SetUpTestCase()
    {
        std::cout << "BroadCastRingFor91093GeneralizationTest set up." << std::endl;
    }
 
    static void TearDownTestCase()
    {
        std::cout << "BroadCastRingFor91093GeneralizationTest tear down." << std::endl;
    }

    virtual void SetUp()
    {
        const ::testing::TestInfo* const test_info = ::testing::UnitTest::GetInstance()->current_test_info();
        std::string caseName = "analysis_result_" + std::string(test_info->test_case_name()) + "_" + std::string(test_info->name());
        Checker::SetDumpFileName(caseName);
    }
 
    virtual void TearDown()
    {
        Checker::SetDumpFileName("analysis_result");
        GlobalMockObject::verify();
        // 这边每个case执行完成需要清理所有的环境变量，如果有新增的环境变量，需要在这个函数中进行清理
        ClearHcclEnv();
    }
protected:
    int GetRootValue(int rankSize, int rootParam){
        switch (rootParam)
        {
        case 0: 
            return 0;
            break;
        case 1: 
            return rankSize - 1;
            break;
        case 2: 
            return rankSize / 2;
            break;
        case 3: 
            return std::min(rankSize / 2 + 1, rankSize -1);
            break;
        default:
            std::cout << "The root node is invalid values." << std::endl;
            return -1;
            break;
        }
    }
};
    
TEST_P(BroadCastRingFor91093GeneralizationTest, SymmetricTopo)
{
    const auto& mytuple = GetParam();
    // 打印一下mytuple当前的参数
    // std::cout << "testOpParam.DataDes.count : " << std::get<0>(mytuple) << std::endl;
    // std::cout << "testOpParam.DataDes.dataType : " << std::get<1>(mytuple) << std::endl;
    // std::cout << "testOpParam.opMode : " << std::get<2>(mytuple) << std::endl;
    // std::cout << "testOpParam.TopoMeta : " << std::get<3>(mytuple)[0]  << " , " << std::get<3>(mytuple)[1] << " , " << std::get<2>(mytuple)[3] << std::endl;
    // std::cout << "testOpParam.root : " << std::get<4>(mytuple) << std::endl;

    RankTable_t ranktable;
    RankTable_For_LLT gen;
    TopoMeta topoMeta;
    gen.GenTopoMeta(topoMeta, std::get<3>(mytuple)[0], std::get<3>(mytuple)[1], std::get<3>(mytuple)[2]);
    gen.GenRankTable(topoMeta, ranktable);
    int rankSize = ranktable.rankNum;

    setenv("HCCL_ALGO", "level0:NA;level1:NHR", 1);
    u64 dataSize = std::get<0>(mytuple);
    if (dataSize >= 5000000008) {
        setenv("HCCL_BUFFSIZE", "4096", 1);
    };
    u64 datacount = dataSize / SIZE_TABLE[std::get<1>(mytuple)];
    
    TestOpParam testOpParam;
    testOpParam.opType = HcclCMDType::HCCL_CMD_BROADCAST;
    testOpParam.tag = "broadcast";
    testOpParam.opMode = std::get<2>(mytuple);
    testOpParam.devtype = DevType::DEV_TYPE_910_93;
    testOpParam.algName = "BroadCastRingFor91093Executor";
    testOpParam.DataDes.count = datacount;
    testOpParam.DataDes.dataType = std::get<1>(mytuple);
    testOpParam.root = GetRootValue(rankSize, std::get<4>(mytuple));
    
    Checker checker;
    HcclResult ret;
    checker.CloseRankMemCheck();
    ret = checker.Check(testOpParam, ranktable, topoMeta);
    EXPECT_EQ(ret, HcclResult::HCCL_SUCCESS);
}

TEST_P(BroadCastRingFor91093GeneralizationTest, AsymmetricTopo)
{
    const auto& mytuple = GetParam();
    // 打印一下mytuple当前的参数
    // std::cout << "testOpParam.DataDes.count : " << std::get<0>(mytuple) << std::endl;
    // std::cout << "testOpParam.DataDes.dataType : " << std::get<1>(mytuple) << std::endl;
    // std::cout << "testOpParam.opMode : " << std::get<2>(mytuple) << std::endl;
    // std::cout << "testOpParam.root : " << std::get<4>(mytuple) << std::endl;

    RankTable_t ranktable;
    RankTable_For_LLT gen;
    TopoMeta topoMeta = {std::get<5>(mytuple)[0]};
    gen.GenRankTable(topoMeta, ranktable);
    int rankSize = ranktable.rankNum;

    setenv("HCCL_ALGO", "level0:NA;level1:NHR", 1);
    u64 dataSize = std::get<0>(mytuple);
    if (dataSize >= 5000000008) {
        setenv("HCCL_BUFFSIZE", "4096", 1);
    };
    u64 datacount = dataSize / SIZE_TABLE[std::get<1>(mytuple)];
    
    TestOpParam testOpParam;
    testOpParam.opType = HcclCMDType::HCCL_CMD_BROADCAST;
    testOpParam.tag = "broadcast";
    testOpParam.opMode = std::get<2>(mytuple);
    testOpParam.devtype = DevType::DEV_TYPE_910_93;
    testOpParam.algName = "BroadCastRingFor91093Executor";
    testOpParam.DataDes.count = datacount;
    testOpParam.DataDes.dataType = std::get<1>(mytuple);
    testOpParam.root = GetRootValue(rankSize, std::get<4>(mytuple));
    
    Checker checker;
    HcclResult ret;
    checker.CloseRankMemCheck();
    ret = checker.Check(testOpParam, ranktable, topoMeta);
    EXPECT_EQ(ret, HcclResult::HCCL_SUCCESS);
}

INSTANTIATE_TEST_CASE_P(BroadCastTest, BroadCastRingFor91093GeneralizationTest, 
    testing::Combine(
        testing::Values(800/*, 1000000008, 5000000008*/), 
        testing::Values(HcclDataType::HCCL_DATA_TYPE_INT32, HcclDataType::HCCL_DATA_TYPE_INT8/*, HcclDataType::HCCL_DATA_TYPE_BFP16, HcclDataType::HCCL_DATA_TYPE_INT64*/),
        testing::Values(OPBASE, OFFLOAD), 
        testing::ValuesIn(std::vector<std::vector<int>>{{1, 1, 8}, {1, 1, 16}, {1, 4, 8}, {1, 4, 16}}),
        testing::Values(0, 1, 2, 3),
        testing::ValuesIn(std::vector<std::vector<std::vector<std::vector<unsigned int>>>>
        {{{{0,1,2,3,4,6,8}, {0,1,2,3,4,6,8}}}})
    )
);